#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""  
@Project : tools-for-internet-security
@File : sec-tools.py
@Author : Misha
@Time : 2025/5/11 12:43  
@脚本说明 : 项目运行主文件
"""
import argparse
import os
import json
from blaster import BlastTools
from mail_operator import send_mail
from file_operator import encrypt_file, decrypt_file, generate_hash, verify_hash


def init():
    # 程序初始化，生成或读取信息记录文件
    if not os.path.exists(".info.json"):
        with open(".info.json", "w", encoding="utf-8"):
            return {}
    else:
        with open(".info.json", "r", encoding="utf-8") as f:
            try:
                return json.load(f)
            except json.JSONDecodeError:
                with open(".info.json", "w", encoding="utf-8") as f1:
                    f1.write("")
                    return {}


def create_parser() -> argparse.Namespace:
    """创建命令行解析器"""
    parser = argparse.ArgumentParser(description="主程序实例对象")
    subparsers = parser.add_subparsers(dest="sub_parsers", description="子目录命令容器")

    # 创建爆破命令容器及选项列表
    blast_parser = subparsers.add_parser("blast", help="密码爆破相关操作")
    # 创建不同爆破协议的互斥选项组
    blast_group = blast_parser.add_mutually_exclusive_group(required=True)
    blast_group.add_argument("-m", "--mysql", action="store_true", help="使用 MySQL 爆破")
    blast_group.add_argument("-r", "--redis", action="store_true", help="使用 Redis 爆破")
    blast_group.add_argument("-s", "--ssh", action="store_true", help="使用 SSH 爆破")
    blast_parser.add_argument("-H", "--host", required=True, help="(必填) 爆破目标主机地址")
    blast_parser.add_argument("-P", "--port", required=True, type=int, help="(必填) 爆破目标主机端口")
    blast_parser.add_argument("-u", "--username", default="root", type=str,
                              help="(选填) 爆破目标用户名，默认为root。Redis不需要指定")
    blast_parser.add_argument("-p", "--password",
                              default="server_password.txt",
                              type=argparse.FileType("r", encoding="utf-8"),
                              help="(选填) 爆破密码本，默认为server_password.txt")
    blast_parser.add_argument("-T", "--threads", default=8, type=int,
                              help="(选填) 爆破线程数，默认为8")
    blast_parser.add_argument("-S", "--sleep", default=0.01, type=float,
                              help="(选填) 爆破线程间休眠时间，默认为0.01s")

    # 创建文件操作容器及选项列表
    file_parser = subparsers.add_parser("file", help="文件相关操作")
    # 创建文件操作互斥选项组
    file_group = file_parser.add_mutually_exclusive_group(required=True)
    file_group.add_argument("-e", "--encrypt", action="store_true", help="加密文件")
    file_group.add_argument("-d", "--decrypt", action="store_true", help="解密文件")
    file_group.add_argument("-g", "--generate", action="store_true", help="生成文件校验值")
    file_group.add_argument("-v", "--verify", action="store_true", help="校验文件")
    file_parser.add_argument("-F", "--file", required=True, type=str, help="(必填) 文件路径")
    file_parser.add_argument("-H", "--hash", type=str, help="(校验文件必填) 校验文件使用的哈希值，支持MD5/SHA1/SHA256")

    # 创建邮件操作容器及选项列表
    mail_parser = subparsers.add_parser("mail", help="发送邮件")
    # 创建发送邮件所需参数选项
    mail_parser.add_argument("-S", "--server", required=True, help="(必填) 邮件服务器地址")
    mail_parser.add_argument("-u", "--user", required=True, help="(必填) 发信人邮箱地址")
    mail_parser.add_argument("-r", "--receiver", required=True, nargs="+", help="(必填) 收信人邮箱地址，可以指定多个")
    mail_parser.add_argument("-p", "--password", default=None, help="(首次使用必填) 邮箱密码")
    mail_parser.add_argument("-f", "--file", default="mail_template.html",
                             type=argparse.FileType("r", encoding="utf-8"),
                             help="(选填) 邮件内容文件路径，默认为mail_template.html")
    mail_parser.add_argument("-P", "--port", default=465, help="(选填) 服务器SMTP端口，默认465")

    # 解析命令行参数
    return parser.parse_args()


def run(arguments: argparse.Namespace):
    """脚本运行"""
    if arguments.sub_parsers == "blast":
        if arguments.mysql:
            BlastTools(method="mysql",
                       host=arguments.host,
                       port=arguments.port,
                       username=arguments.username,
                       password_list=arguments.password.readlines(),
                       max_threads=arguments.threads,
                       sleep=arguments.sleep).run()
        elif arguments.redis:
            BlastTools(method="redis",
                       host=arguments.host,
                       port=arguments.port,
                       username='',
                       password_list=arguments.password.readlines(),
                       max_threads=arguments.threads,
                       sleep=arguments.sleep).run()
        elif arguments.ssh:
            if arguments.threads > 8:
                ssh_choice = input("线程数过高时，可能会导致SSH爆破失败，是否继续？(y/n)\n")
                if ssh_choice == "y":
                    BlastTools(method="ssh",
                               host=arguments.host,
                               port=arguments.port,
                               username=arguments.username,
                               password_list=arguments.password.readlines(),
                               max_threads=arguments.threads).run()
                else:
                    return
            else:
                BlastTools(method="ssh",
                           host=arguments.host,
                           port=arguments.port,
                           username=arguments.username,
                           password_list=arguments.password.readlines(),
                           max_threads=arguments.threads).run()
    elif arguments.sub_parsers == "file":
        if not os.path.exists(arguments.file):
            print("文件不存在！")
            return
        if arguments.encrypt:
            encrypt_file(file_path=arguments.file)
        elif arguments.decrypt:
            decrypt_file(file_path=arguments.file)
        elif arguments.generate:
            generate_hash(file_path=arguments.file)
        elif arguments.verify:
            if not arguments.hash:
                print("请给定校验文件使用的哈希值！")
                return
            verify_hash(file_path=arguments.file, hash_value=arguments.hash)
    elif arguments.sub_parsers == "mail":
        send_mail(info_dict=info, mail_content=arguments.file, receiver_list=arguments.receiver,
                  smtp_server=arguments.server, smtp_port=arguments.port,
                  sender_email=arguments.user, sender_password=arguments.password)
    else:
        print("未知命令，请重新输入！")


if __name__ == '__main__':
    info: dict = init()
    args: argparse.Namespace = create_parser()
    run(args)
